import os
import torch
import SimpleITK as sitk

# 图像地址
input_image_path = 'C:/SYBdataset/dataset/qutougu/brain'
input_image_path_files = [os.path.join(input_image_path, f) for f in os.listdir(input_image_path) if f.endswith('.nii.gz')]

# 标签地址
input_mask_path = 'C:/SYBdataset/dataset/AffinedManualSegImageNIfTI'
input_mask_path_files = [os.path.join(input_mask_path, f) for f in os.listdir(input_mask_path) if f.endswith('.nii.gz')]

# 保存图像的地址
output_image_path = 'C:/SYBdataset/dataset/image_registration/RawImageNIfTI'
output_mask_path = 'C:/SYBdataset/dataset/image_registration/AffinedManualSegImageNIfTI'

# 用于配准的图像地址
image_s001 = sitk.ReadImage('C:/SYBdataset/dataset/qutougu/brain/s001.nii.gz')
Origin = image_s001.GetOrigin()  # 获取原点坐标
Spacing = image_s001.GetSpacing()  # 获取体素值
Direction = image_s001.GetDirection()  # 获取方向余弦矩阵
for i in range(len(input_image_path_files)):
    # 读取对应的图像和标签
    img = sitk.ReadImage(input_image_path_files[i])
    mak = sitk.ReadImage(input_mask_path_files[i])
    # 将图像进行匹配
    img.SetOrigin(Origin)
    img.SetSpacing(Spacing)
    img.SetDirection(Direction)

    # 将标签进行匹配
    mak.SetOrigin(Origin)
    mak.SetSpacing(Spacing)
    mak.SetDirection(Direction)

    # 保存图像和标签
    sitk.WriteImage(img, os.path.join(output_image_path, f'{i}.nii.gz'))
    sitk.WriteImage(mak, os.path.join(output_mask_path, f'{i}.nii.gz'))

